# -*- utf-8 -*-
import tensorflow as tf

#配置网络参数
input_node=784
output_node=10

image_size=28
num_channels=1
num_labels=10

#第一层卷积层的尺寸和深度
conv1_deep=32   #过滤器数量
conv1_size=5   #过滤器f尺寸
#第二层
conv2_deep=64
conv2_size=5
#全连接层的节点个数
fc_size=512

#定义卷积神经网络的前向传播过程，这里新增了新的参数train，用于区分训练过程和测试过程。在这个程序中使用dropout方法，
#来进一步提升模型的可靠性并防止过拟合，dropout只在训练时使用。
def inference(input_tensor,train,regularizer):
    '''
    :param input_tensor: 网络输入
    :param train: 是训练还是测试
    :param regularizer: 正则化器，有的话则进行正则化损失的计算
    :return: 前向传播结果，即第六层的输出
    这里网络结构和标准的LeNet-5有多不同，但基本类似.
    '''
    print(input_tensor.shape)
    #第一层卷积层的前向传播,  (batch,28*28*1)*(5*5*1*32)=(batch,28*28*32)  same卷积则会填充，使得前后尺寸不变
    with tf.variable_scope('layer1-conv1'):
        #mnist数据集图像为28*28*1
        conv1_weights=tf.get_variable(
            name='weight',shape=[conv1_size,conv1_size,num_channels,conv1_deep],initializer=tf.truncated_normal_initializer(stddev=0.1)
        )
        conv1_biases=tf.get_variable(
            'bias',[conv1_deep],initializer=tf.constant_initializer(0.0)
        )
        conv1=tf.nn.conv2d(
            input_tensor,conv1_weights,strides=[1,1,1,1],padding="SAME"
        )
        relu1=tf.nn.relu(tf.nn.bias_add(conv1,conv1_biases))

    #第二层池化层的前向传播,   (batch,28*28*32)变为(batch,14*14*32)
    with tf.variable_scope('layer2-pool1'):
        pool1=tf.nn.max_pool(
            relu1,ksize=[1,3,3,1],strides=[1,2,2,1],padding="SAME"
        )

    #第三层卷积层的前向传播,  (batch,14*14*32)*(5*5*32*64)=(batch,14*14*64)
    with tf.variable_scope('layer3-conv2'):
        conv2_weights=tf.get_variable(
            'weight',[conv2_size,conv2_size,conv1_deep,conv2_deep],initializer=tf.truncated_normal_initializer(stddev=0.1)
        )
        conv2_biases=tf.get_variable(
            'bias',[conv2_deep],initializer=tf.constant_initializer(0.0)
        )
        conv2=tf.nn.conv2d(
            pool1,conv2_weights,strides=[1,1,1,1],padding="SAME"
        )
        relu2=tf.nn.relu(tf.nn.bias_add(conv2,conv2_biases))

    #第四层池化层的前向传播,   (batch,14*14*64)变为(batch,7*7*64)
    with tf.variable_scope('layer4-pool2'):
        pool2=tf.nn.max_pool(
            relu2,ksize=[1,2,2,1],strides=[1,2,2,1],padding="SAME"
        )

    #第五层全连接层的前向传播，(batch,7*7*64)变为（batch,3136）,然后和(3136,512)进行全连接运算，得到(batch,512)
    # 因此要先将第四层的结果矩阵7*7*64进行平铺展开为向量然后送入
    #注意，因为实际每一层输入输出都为一个batch矩阵，所以这里得到的维度也包含了一个batch中数据的个数
    #(batch,7,7,64)
    pool_shape=pool2.get_shape().as_list()
    print(pool_shape)
    nodes=pool_shape[1]*pool_shape[2]*pool_shape[3]   #7*7*64
    #通过tf.reshape()函数将第四层的输出变成一个batch的向量
    reshaped=tf.reshape(pool2,[pool_shape[0],nodes])
    #dropout会在训练时随即将部分节点的输出改为0，以此来避免过拟合，dropout一般只在全连接层使用而不再卷积层和池化层使用
    with tf.variable_scope('layer5-fc1'):
        fc1_weights=tf.get_variable(
            'weight',[nodes,fc_size],initializer=tf.truncated_normal_initializer(stddev=0.1)
        )
        if regularizer!=None:
            tf.add_to_collection('losses',regularizer(fc1_weights))
        fc1_biases=tf.get_variable(
            'bias',[fc_size],initializer=tf.constant_initializer(0.1)
        )
        fc1=tf.nn.relu(tf.matmul(reshaped,fc1_weights)+fc1_biases)
        if train:fc1=tf.nn.dropout(fc1,0.5)

    #第六层全连接层的前向传播,  (batch,512)和(512,10)=(batch,10)
    with tf.variable_scope('layer6-fc2'):
        fc2_weights=tf.get_variable(
            'weight',[fc_size,num_labels],initializer=tf.truncated_normal_initializer(stddev=0.1)
        )
        fc2_biases=tf.get_variable(
            'biases',[num_labels],initializer=tf.constant_initializer(0.1)
        )
        #一般对于分类，最后的返回用logit表示，即logit=fc2
        fc2=tf.nn.relu(tf.matmul(fc1,fc2_weights)+fc2_biases)

    return fc2

